src.clean

 1import re
 2import html
 3import emoji
 4import logging
 5from bs4 import BeautifulSoup
 6from nltk.corpus import words, stopwords
 7from nltk.stem import WordNetLemmatizer
 8
 9logger = logging.getLogger("yt_pipeline")
10
11# Preload NLTK resources for performance
12wordset = set(words.words())
13stopwordset = set(stopwords.words('english'))
14lemmatizer = WordNetLemmatizer()
15
16def clean_text(text):
17    """
18    Cleans raw comment text by removing HTML tags, unescaping HTML entities,
19    demojizing emojis, and stripping user mentions.
20
21    Args:
22        text (str): The raw text to clean.
23
24    Returns:
25        str: A cleaned version of the input text.
26    """
27    try:
28        text = re.sub(r'<!--(.*?)-->', '', text)
29        text = re.sub(r'<i>(.*?)</i>', r'\1', text)
30        text = re.sub(r'<a(.*?)>(.*?)</a>', r'\2', text)
31        text = html.unescape(text)
32        return emoji.demojize(BeautifulSoup(text, 'html.parser').get_text(separator='\n'))
33    except Exception as e:
34        logger.warning(f"[clean_text] Failed to clean text: {text[:60]}... | Error: {e}")
35        return ""
36
37def clean_comments(df):
38    """
39    Applies text cleaning to all rows in the DataFrame's 'textDisplay' column.
40
41    Args:
42        df (pd.DataFrame): A DataFrame containing a 'textDisplay' column.
43
44    Returns:
45        pd.DataFrame: The DataFrame with cleaned text.
46    """
47    try:
48        df['textDisplay'] = df['textDisplay'].astype(str)
49        df['textDisplay'] = df['textDisplay'].str.replace(r'@\w+', '', regex=True)
50        df['textDisplay'] = df['textDisplay'].apply(clean_text)
51        logger.info("[clean_comments] Finished cleaning text.")
52    except Exception as e:
53        logger.error("[clean_comments] Error during text cleaning", exc_info=True)
54    return df
55
56def tokenize_comments(df):
57    """
58    Tokenizes cleaned comment text into lemmatized words, removing stopwords
59    and non-English words.
60
61    Args:
62        df (pd.DataFrame): A DataFrame with a 'textDisplay' column containing cleaned text.
63
64    Returns:
65        pd.DataFrame: The DataFrame with a new 'wordtoken' column containing token lists.
66    """
67    def tokenizer(sentence):
68        try:
69            return [
70                lemmatizer.lemmatize(word.lower())
71                for word in sentence.split()
72                if word.lower() in wordset and word.lower() not in stopwordset
73            ]
74        except Exception as e:
75            logger.warning(f"[tokenizer] Failed to tokenize: {sentence[:60]}... | Error: {e}")
76            return []
77
78    try:
79        df['wordtoken'] = df['textDisplay'].apply(tokenizer)
80        empty_tokens = df[df['wordtoken'].apply(lambda x: len(x) == 0)]
81
82        if not empty_tokens.empty:
83            logger.warning(f"[tokenize_comments] {len(empty_tokens)} rows resulted in empty tokens.")
84            logger.debug(empty_tokens[['textDisplay']].head(3).to_string(index=False))
85        else:
86            logger.info("[tokenize_comments] All rows successfully tokenized.")
87    except Exception as e:
88        logger.error("[tokenize_comments] Error applying tokenizer", exc_info=True)
89
90    return df